package Service

import (
	"forward-core/Utils"
	"github.com/astaxie/beego/logs"
	"net"
	"os"
	"time"
)

type ForWardClient struct {
	srcConn        net.Conn
	destAddr       string
	dispatchAddrs  string
	closedCallBack func(srcConn net.Conn, destConn net.Conn)

	dispatchMap map[string]*myWriter
	destConn    net.Conn
	destChan    chan []byte
	isRunning   bool
}

func NewForWardClient(SrcConn net.Conn, DestAddr string, DispatchAddrs string,
	ClosedCallBack func(srcConn net.Conn, destConn net.Conn)) *ForWardClient {
	f := &ForWardClient{
		srcConn:        SrcConn,
		destAddr:       DestAddr,
		dispatchAddrs:  DispatchAddrs,
		closedCallBack: ClosedCallBack,
	}
	f.destChan = make(chan []byte, 100)
	f.dispatchMap = make(map[string]*myWriter)
	f.isRunning = true
	return f
}

var f *os.File

func (_self *ForWardClient) DispatchData() {
	//// 开始CPU剖析
	//f, err := os.Create("cpu.prof")
	//if err != nil {
	//	logs.Error(err)
	//}
	//pprof.StartCPUProfile(f)

	//启动其他数据分发连接
	if Utils.IsNotEmpty(_self.dispatchAddrs) {
		dispatchTargets := Utils.Split(_self.dispatchAddrs, ";")
		for _, v := range dispatchTargets {
			aWriter := &myWriter{
				destAddr:  v,
				bufChan:   make(chan []byte, 100),
				isRunning: true,
			}
			aWriter.start()
			_self.dispatchMap[v] = aWriter
		}
	}
	//读取SrcConn写入chan
	go func() {
		var buffer = make([]byte, 1024)
		for _self.isRunning {
			n, err := _self.srcConn.Read(buffer)
			if err != nil {
				logs.Error(err)
				_self.StopForward()
				return
			}
			if n > 0 {
				tmp := make([]byte, n)
				copy(tmp, buffer[0:n])
				WriteChWithSelect(_self.destChan, tmp)
				for _, v := range _self.dispatchMap {
					WriteChWithSelect(v.bufChan, tmp)
				}
			}
		}
	}()
	//读取chan写入destConn目标
	go func() {
		for _self.isRunning {
			var err error = nil
			_self.destConn = nil
			_self.destConn, err = net.DialTimeout("tcp", _self.destAddr, 10*time.Second)

			if err != nil {
				logs.Debug("连接主目标超时: %s, %s", _self.destAddr, err.Error())
				time.Sleep(time.Millisecond * 1000 * 10)
				continue
			}

			for v := range _self.destChan {
				if !_self.isRunning {
					break
				}
				_, err = _self.destConn.Write(v)
				if err != nil {
					logs.Error("主目标转发数据: ", err.Error())
					break
				}
				logs.Debug("主目标转发数据: ", string(v))
				time.Sleep(time.Second * 0)
			}
			CloseConn(_self.destConn)
			logs.Debug("转发主目标出错：", _self.destAddr)
		}
	}()
	//读取destConn，写入srcConn,(出错重新连接在上一个函数)
	go func() {
		buf := make([]byte, 1024)
		for _self.isRunning {
			if _self.destConn != nil && _self.srcConn != nil {
				func() {
					defer func() {
						if r := recover(); r != nil {
							logs.Error("Recovered from panic:", r) // 如果关闭连接时发生panic，恢复并打印信息
						}
					}()
					//err2 := _self.destConn.SetReadDeadline(time.Now().Add(time.Second * 10))
					//if err2 != nil {
					//	CloseConn(_self.destConn)
					//	time.Sleep(time.Millisecond * 100)
					//	return
					//}
					n, _ := _self.destConn.Read(buf)
					if n <= 0 {
						//CloseConn(_self.destConn)
						time.Sleep(time.Millisecond * 100)
						//fmt.Println("___________________read data length is zero")
						return
					} else {
						_self.srcConn.Write(buf[:n])
					}
				}()
			} else {
				time.Sleep(time.Millisecond * 100)
			}
		}
	}()
}

func (_self *ForWardClient) StopForward() {
	logs.Debug("关闭一个连接：", _self.srcConn.RemoteAddr(), " on ", _self.srcConn.LocalAddr())
	_self.isRunning = false
	if Utils.IsNotEmpty(_self.dispatchAddrs) {
		dispatchTargets := Utils.Split(_self.dispatchAddrs, ";")
		for _, v := range dispatchTargets {
			writer := _self.dispatchMap[v]
			if writer != nil {
				writer.stop()
			}

		}
	}
	CloseChan(_self.destChan)
	CloseConn(_self.srcConn)
	CloseConn(_self.destConn)
	_self.closedCallBack(_self.srcConn, _self.destConn)
	logs.Debug("all connection closed!")

	//// 停止CPU剖析并关闭文件
	//pprof.StopCPUProfile()
	//f.Close()
}
